{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_pred correct/all (2444/2921); num_gt correct/all (2444/3068)\n",
      "precision: 83.669976%\n",
      "recall: 79.661017%\n",
      "f-measure: 81.616297%\n"
     ]
    }
   ],
   "source": [
    "from argparse import ArgumentParser\n",
    "import mmcv\n",
    "import numpy as np\n",
    "import os\n",
    "import os.path as osp\n",
    "import numpy as np\n",
    "from shapely.geometry import Polygon, LineString, MultiLineString, GeometryCollection\n",
    "from shapely.algorithms.polylabel import polylabel\n",
    "import math\n",
    "import cv2\n",
    "import json\n",
    "import pickle\n",
    "\n",
    "def load_ann(filename):\n",
    "    with open(filename) as f:\n",
    "        texts = f.readlines()\n",
    "    num_objs = len(texts)\n",
    "    polygon_list = []\n",
    "    for i in range(num_objs):\n",
    "        try:\n",
    "            box = texts[i].split(',')\n",
    "            box = np.array(box, dtype=np.float32)\n",
    "            box[4::2] += box[0]\n",
    "            box[5::2] += box[1]\n",
    "            box = box[4:]\n",
    "            hard = 0\n",
    "            polygon = Polygon(box.reshape((-1, 2)))\n",
    "        except Exception as err:\n",
    "            print(filename, texts[i], err)\n",
    "        polygon_list.append([polygon, hard])\n",
    "    return polygon_list\n",
    "\n",
    "def load_pred(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        array = pickle.load(f)\n",
    "    pred_polygon_list = []\n",
    "    for ix, pred_per_sample in enumerate(array):\n",
    "        polygon_list = []\n",
    "        for pred_one_out in pred_per_sample:\n",
    "            coords = pred_one_out[:-1]\n",
    "            score = float(pred_one_out[-1])\n",
    "#             coords = pred_one_out\n",
    "#             score = 1\n",
    "            if score < 0.95:\n",
    "                continue\n",
    "            polygon = Polygon(coords.reshape((-1, 2)))\n",
    "#             if not polygon.is_valid:\n",
    "#                 print(ix, polygon)\n",
    "            polygon_list.append(polygon)\n",
    "        pred_polygon_list.append(polygon_list)\n",
    "    return pred_polygon_list\n",
    "\n",
    "def polygon_iou(poly1, poly2):\n",
    "    \"\"\"\n",
    "    Intersection over union between two shapely polygons.\n",
    "    \"\"\"\n",
    "    iou = 0.0\n",
    "    if poly1.intersects(poly2):\n",
    "        try:\n",
    "            inter_area = poly1.intersection(poly2).area\n",
    "            union_area = poly1.area + poly2.area - inter_area\n",
    "            iou = float(inter_area) / union_area\n",
    "        except Exception as err:\n",
    "            iou = 0.5\n",
    "#             print(poly1)\n",
    "    return iou\n",
    "\n",
    "def evaluation(pred_polygon_list, gt_polygon_list):\n",
    "    assert len(pred_polygon_list)==len(gt_polygon_list), \"sample numbers should be same\"\n",
    "    num_gt = 0\n",
    "    num_pred = 0\n",
    "    num_gt_correct = 0\n",
    "    num_pred_correct = 0\n",
    "    for sample_idx in range(len(pred_polygon_list)):\n",
    "        gt_polygons = gt_polygon_list[sample_idx]\n",
    "        pred_polygons = pred_polygon_list[sample_idx]\n",
    "        num_gt += len(gt_polygons)\n",
    "        num_pred += len(pred_polygons)\n",
    "        for gt_polygon, hard in gt_polygons:\n",
    "            num_gt -= 1 if hard else 0\n",
    "            gt_correct = False\n",
    "            for pred_polygon in pred_polygons:\n",
    "                iou = polygon_iou(gt_polygon, pred_polygon)\n",
    "                if iou >= 0.5:\n",
    "                    num_pred_correct += 1 if not hard else 0\n",
    "                    num_pred -= 1 if hard else 0\n",
    "                    gt_correct = True\n",
    "            num_gt_correct += 1 if (gt_correct and not hard) else 0\n",
    "    p = 1.0 * num_pred_correct / max(num_pred, 1e-10)\n",
    "    r = 1.0 * num_gt_correct / max(num_gt, 1e-10)\n",
    "    f = 2.0 * p * r / max(p + r, 1e-10)\n",
    "    print(\"num_pred correct/all (%d/%d); num_gt correct/all (%d/%d)\"%(num_pred_correct, num_pred, num_gt_correct, num_gt))\n",
    "    print(\"precision: %f%%\"%(p*100.0))\n",
    "    print(\"recall: %f%%\"%(r*100.0))\n",
    "    print(\"f-measure: %f%%\"%(f*100.0))\n",
    "        \n",
    "# load pred pickle\n",
    "pickle_file = '../../experiments/CTW_cheby/CTW_cheby_08051049.pkl'\n",
    "pred_polygon_list = load_pred(pickle_file)\n",
    "# load gt ann\n",
    "img_prefix = '../../data/CTW1500/'\n",
    "test_file = os.path.join(img_prefix, 'ImageSets/Main/test.txt')\n",
    "img_ids = mmcv.list_from_file(test_file)\n",
    "gt_polygon_list = []\n",
    "for img_id in img_ids:\n",
    "    ann_file = os.path.join(img_prefix, 'Annotations', img_id + '.txt')\n",
    "    gt_polygon_list.append(load_ann(ann_file))\n",
    "# evaluation\n",
    "result = evaluation(pred_polygon_list, gt_polygon_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
